import sys, os
from random import shuffle

def writeFile(bags, n, filename):
# escreve o arquivo de teste ou treinamento no formato da libsvm

    fid = open(filename, 'a')
    for bag in bags:
        fid.write(str(n))    # escreve label

        # arquivo esparso: so escreve os valores diferentes de zero    
        for i in range(len(bag)):
            if float(bag[i]) != 0:
                fid.write(" " + str(i+1) + ":" + bag[i])
        fid.write('\n')

    fid.close()


def separateSets(bags, classe, n, trainFile, testFile):
# separa conjunto de treinamento com n amostras aleatorias de uma classe da base, e o restante das amostras sera do conjunto de teste

    shuffle(bags)    # sort list of bags

    writeFile(bags[:n], classe, trainFile) # write train file
    writeFile(bags[n:], classe, testFile)  # write test file


def readBags(filename):
# le o conjunto de bags, separando-as por classe

    bags = dict()  # dicionario cujas chaves sao nomes de classes da base (diretorio das imagens)

    fid = open(filename)
    lines = fid.readlines()
    fid.close()

    # cada classe(chave) do dicionario 'bags' recebera uma lista de bags (cada bag esta associada a uma imagem da base)
    for line in lines:
        sl = line.split()
        classe, img = sl[0].split('/')
        if bags.has_key(classe):
            bags[classe].append(sl[1:])
        else:
            bags[classe] = [sl[1:]]
        
    return bags


def main():

    bagsFile = sys.argv[1] # arquivo com as bags geradas
    n = int(sys.argv[2])  # numero de amostras de treinamento

    trainFile = sys.argv[3] # arquivo com amostras de treino
    testFile = sys.argv[4] # arquivos com amostras de teste

    bags = readBags(bagsFile)
    classe=1   # label da classe
    for k in bags.keys():
	    # separa as amostras de treinamento e teste para cada classe 
        separateSets(bags[k], classe, n, trainFile, testFile)
	classe+=1


if __name__ == "__main__":
    main()
